package com.jzz.algorithm;

import java.util.Arrays;

public class DesAlgorithm implements Algorithm {

    /** 初始置换函数 */
    private static final int[] IP = new int[] {
            58,50,42,34,26,18,10, 2,
            60,52,44,36,28,20,12, 4,
            62,54,46,38,30,22,14, 6,
            64,56,48,40,32,24,16, 8,
            57,49,41,33,25,17, 9, 1,
            59,51,43,35,27,19,11, 3,
            61,53,45,37,29,21,13, 5,
            63,55,47,39,31,23,15, 7
    };

    /** 逆初始置换函数 */
    private static final int[] RIP = new int[] {
            40, 8,48,16,56,24,64,32,
            39, 7,47,15,55,23,63,31,
            38, 6,46,14,54,22,62,30,
            37, 5,45,13,53,21,61,29,
            36, 4,44,12,52,20,60,28,
            35, 3,43,11,51,19,59,27,
            34, 2,42,10,50,18,58,26,
            33, 1,41, 9,49,17,57,25,
    };

    /** 移动位数 */
    private static final int[] LFT = {1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1};

    /** 扩展置换 */
    private static final int[] E = {
            32, 1, 2, 3, 4, 5,
             4, 5, 6, 7, 8, 9,
             8, 9,10,11,12,13,
            12,13,14,15,16,17,
            16,17,18,19,20,21,
            20,21,22,23,24,25,
            24,25,26,27,28,29,
            28,29,30,31,32, 1
    };

    /** 密钥置换函数 */
    private static final int[] PC1 = {
            57,49,41,33,25,17, 9,
             1,58,50,42,34,26,18,
            10, 2,59,51,43,35,27,
            19,11, 3,60,52,44,36,
            63,55,47,39,31,23,15,
             7,62,54,46,38,30,22,
            14, 6,61,53,45,37,29,
            21,13, 5,28,20,12, 4
    };

    /** 压缩置换 */
    private static final int[] PC2 = {
            14,17,11,24, 1, 5, 3,28,
            15, 6,21,10,23,19,12, 4,
            26, 8,16, 7,27,20,13, 2,
            41,52,31,37,47,55,30,40,
            51,45,33,48,44,49,39,56,
            34,53,46,42,50,36,29,32
    };

    /** P盒置换 */
    private static final int[] P = {
            16, 7,20,21,29,12,28,17,
             1,15,23,26, 5,18,31,10,
             2, 8,24,14,32,27, 3, 9,
            19,13,30, 6,22,11, 4,25
    };

    /** S盒 */
    private static final int[][][] S_BOX = {
            // S-BOX 1
            {
                    {14, 4,13, 1, 2,15,11, 8, 3,10, 6,12, 5, 9, 0, 7},
                    { 0,15, 7, 4,14, 2,13, 1,10, 6,12,11, 9, 5, 3, 8},
                    { 4, 1,14, 8,13, 6, 2,11,15,12, 9, 7, 3,10, 5, 0},
                    {15,12, 8, 2, 4, 9, 1, 7, 5,11, 3,14,10, 0, 6,13}
            },
            // S-BOX 2
            {
                    {15, 1, 8,14, 6,11, 3, 4, 9, 7, 2,13,12, 0, 5,10},
                    { 3,13, 4, 7,15, 2, 8,14,12, 0, 1,10, 6, 9,11, 5},
                    { 0,14, 7,11,10, 4,13, 1, 5, 8,12, 6, 9, 3, 2,15},
                    {13, 8,10, 1, 3,15, 4, 2,11, 6, 7,12, 0, 5,14, 9}
            },
            // S-BOX 3
            {
                    {10, 0, 9,14, 6, 3,15, 5, 1,13,12, 7,11, 4, 2, 8},
                    {13, 7, 0, 9, 3, 4, 6,10, 2, 8, 5,14,12,11,15, 1},
                    {13, 6, 4, 9, 8,15, 3, 0,11, 1, 2,12, 5,10,14, 7},
                    { 1,10,13, 0, 6, 9, 8, 7, 4,15,14, 3,11, 5, 2,12}
            },
            // S-BOX 4
            {
                    { 7,13,14, 3, 0, 6, 9,10, 1, 2, 8, 5,11,12, 4,15},
                    {13, 8,11, 5, 6,15, 0, 3, 4, 7, 2,12, 1,10,14, 9},
                    {10, 6, 9, 0,12,11, 7,13,15, 1, 3,14, 5, 2, 8, 4},
                    { 3,15, 0, 6,10, 1,13, 8, 9, 4, 5,11,12, 7, 2,14}
            },
            // S-BOX 5
            {
                    { 2,12, 4, 1, 7,10,11, 6, 8, 5, 3,15,13, 0,14, 9},
                    {14,11, 2,12, 4, 7,13, 1, 5, 0,15,10, 3, 9, 8, 6},
                    { 4, 2, 1,11,10,13, 7, 8,15, 9,12, 5, 6, 3, 0,14},
                    {11, 8,12, 7, 1,14, 2,13, 6,15, 0, 9,10, 4, 5, 3}
            },
            // S-BOX 6
            {
                    {12, 1,10,15, 9, 2, 6, 8, 0,13, 3, 4,14, 7, 5,11},
                    {10,15, 4, 2, 7,12, 9, 5, 6, 1,13,14, 0,11, 3, 8},
                    { 9,14,15, 5, 2, 8,12, 3, 7, 0, 4,10, 1,13,11, 6},
                    { 4, 3, 2,12, 9, 5,15,10,11,14, 1, 7, 6, 0, 8,13}
            },
            // S-BOX 7
            {
                    { 4,11, 2,14,15, 0, 8,13, 3,12, 9, 7, 5,10, 6, 1},
                    {13, 0,11, 7, 4, 9, 1,10,14, 3, 5,12, 2,15, 8, 6},
                    { 1, 4,11,13,12, 3, 7,14,10,15, 6, 8, 0, 5, 9, 2},
                    { 6,11,13, 8, 1, 4,10, 7, 9, 5, 0,15,14, 2, 3,12}
            },
            // S-BOX 8
            {
                    {13, 2, 8, 4, 6,15,11, 1,10, 9, 3,14, 5, 0,12, 7},
                    { 1,15,13, 8,10, 3, 7, 4,12, 5, 6,11, 0,14, 9, 2},
                    { 7,11, 4, 1, 9,12,14, 2, 0, 6,10,13,15, 3, 5, 8},
                    { 2, 1,14, 7, 4,10, 8,13,15,12, 9, 0, 3, 5, 6,11}
            }
    };

    private int[][] keys;

    @Override
    public Algorithm setKey(byte[] key) {
        this.keys = generateSubKeys(key);
        return this;
    }

    @Override
    public byte[] encrypt(byte[] m) {
        return messageDivision(m, false);
    }

    @Override
    public byte[] decrypt(byte[] c) {
        return messageDivision(c, true);
    }

    /**
     * 64位明文分组
     */
    private byte[] messageDivision(byte[] data, boolean decrypt) {
        byte[] copyOfData = data.clone();
        int groupCount = data.length / 8;
        int restCount = 8 - (data.length - groupCount * 8);
        byte[] padding;

        if (restCount < 8 || !decrypt) {
            padding = new byte[data.length + restCount];
            System.arraycopy(copyOfData, 0, padding, 0, data.length);
            for (int i = 0; i < restCount; i++) {
                // ansix923 填充
                padding[data.length + i] = (byte) restCount;
            }
        } else {
            padding = copyOfData.clone();
        }
        groupCount = padding.length / 8;

        byte[] piece = new byte[8];
        byte[] result = new byte[padding.length];
        for (int i = 0; i < groupCount; i++) {
            System.arraycopy(padding, i * 8, piece, 0, 8);
            System.arraycopy(cryptoFunction(piece, decrypt), 0, result, i * 8, 8);
        }

        if (decrypt) {
            byte paddingLength = result[result.length - 1];
            byte[] decResult = new byte[result.length - paddingLength];
            System.arraycopy(result, 0, decResult, 0, result.length - paddingLength);
            return decResult;
        }

        return result;
    }

    private byte[] cryptoFunction(byte[] data, boolean decrypt) {
        // 将64位数据进行IP
        int[] initData = new int[64];
        for (int i = 0; i < data.length; i++) {
            String bin = intTo8BitBinary(data[i]);
            for (int j = 0; j < 8; j++) {
                initData[i * 8 + j] = bin.charAt(j) - '0';
            }
        }

        int[] permuteData = new int[64];
        for (int i = 0; i < 64; i++) {
            permuteData[i] = initData[IP[i] - 1];
        }

        if (decrypt) {
            for (int i = 15; i >= 0; i--) {
                l(permuteData, i, decrypt, keys[i]);
            }
        } else {
            for (int i = 0; i < 16; i++) {
                l(permuteData, i, decrypt, keys[i]);
            }
        }

        int[] cipherBits = new int[64];
        for (int i = 0; i < RIP.length; i++) {
            cipherBits[i] = permuteData[RIP[i] - 1];
        }

        byte[] cipher = new byte[8];
        for (int i = 0; i < 8; i++) {
            cipher[i] = (byte) (
                    (cipherBits[i * 8] << 7) +
                            (cipherBits[i * 8 + 1] << 6) +
                            (cipherBits[i * 8 + 2] << 5) +
                            (cipherBits[i * 8 + 3] << 4) +
                            (cipherBits[i * 8 + 4] << 3) +
                            (cipherBits[i * 8 + 5] << 2) +
                            (cipherBits[i * 8 + 6] << 1) +
                            (cipherBits[i * 8 + 7])
            );
        }

        return cipher;
    }

    private void l(int[] m, int index, boolean decrypt, int[] key) {
        int[] l0 = new int[32];
        int[] r0 = new int[32];
        int[] l1;
        int[] r1 = new int[32];
        int[] f;
        System.arraycopy(m, 0, l0, 0, 32);
        System.arraycopy(m, 32, r0, 0, 32);

        l1 = r0;
        f = function(r0, key);

        for (int i = 0; i < 32; i++) {
            r1[i] = l0[i] ^ f[i];
            if ((decrypt && (index == 0)) || ((!decrypt) && (index == 15))) {
                m[i] = r1[i];
                m[i + 32] = l1[i];
            }
            else {
                m[i] = l1[i];
                m[i + 32] = r1[i];
            }
        }
    }

    private int[] function(int[] rData, int[] key) {
        int[] result = new int[32];
        int[] extKey = new int[48];
        for (int i = 0; i < E.length; i++) {
            extKey[i] = rData[E[i] - 1] ^ key[i];
        }

        // S-Box
        int[][] s = new int[8][6];
        int[] permutedS = new int[32];
        for (int i = 0; i < 8; i++) {
            System.arraycopy(extKey, i * 6, s[i], 0, 6);

            // 横坐标
            int r = (s[i][0] << 1) + s[i][5];

            // 纵坐标
            int c = (s[i][1] << 3) + (s[i][2] << 2) + (s[i][3] << 1) + s[i][4];

            String str = intToBitBinary(S_BOX[i][r][c], 4);

            for (int j = 0; j < 4; j++) {
                permutedS[i * 4 + j] = str.charAt(j) - '0';
            }
        }

        // P-Box
        for (int i = 0; i < P.length; i++) {
            result[i] = permutedS[P[i] - 1];
        }

        return result;
    }

    /**
     * 生成子密钥
     *
     * @return 二进制子密钥
     */
    private int[][] generateSubKeys(byte[] key) {
        int[][] subKeys = new int[16][48];
        if (key.length == 0) {
            return subKeys;
        }

        // 保证密钥8字节
        byte[] trueKey = new byte[8];
        for (int i = 0; i < 8; i++) {
            trueKey[i] = key[i % key.length];
        }

        // 将key转换为64位二进制位
        int[] keyBits = new int[64];
        for (int i = 0; i < 8; i++) {
            String keyStr = intTo8BitBinary(trueKey[i]);

            for (int j = 0; j < 8; j++) {
                int bit = keyStr.charAt(j) - '0';
                if (bit < 0 || bit > 1) {
                    System.out.println("Error: Key generation error.");
                    return subKeys;
                }
                keyBits[i * 8 + j] = bit;
            }
        }

        // PC1 密钥置换函数压缩，最终长度56位
        int[] permutedKey = new int[56];
        for (int i = 0; i < PC1.length; i++) {
            permutedKey[i] = keyBits[PC1[i] - 1];
        }

        // 子密钥
        int[] c0 = new int[28];
        int[] d0 = new int[28];
        System.arraycopy(permutedKey, 0, c0, 0, 28);
        System.arraycopy(permutedKey, 28, d0, 0, 28);

        for (int i = 0; i < 16; i++) {
            int[] c1 = new int[28];
            int[] d1 = new int[28];

            if (LFT[i] == 1) {
                System.arraycopy(c0, 1, c1, 0, 27); c1[27] = c0[0];
                System.arraycopy(d0, 1, d1, 0, 27); d1[27] = d0[0];
            } else {
                System.arraycopy(c0, 2, c1, 0, 26); c1[26] = c0[0]; c1[27] = c0[1];
                System.arraycopy(d0, 2, d1, 0, 26); d1[26] = d0[0]; d1[27] = d0[1];
            }

            int[] tmp = new int[56];
            System.arraycopy(c1, 0, tmp, 0, 28);
            System.arraycopy(d1, 0, tmp, 28, 28);

            // PC2 压缩置换
            for (int j = 0; j < PC2.length; j++) {
                subKeys[i][j] = tmp[PC2[j] - 1];
            }

            c0 = c1;
            d0 = d1;
        }

        return subKeys;
    }

    private String intTo8BitBinary(int v) {
        StringBuilder sb = new StringBuilder(Integer.toBinaryString(v & 0xFF));
        while (sb.length() % 8 != 0) {
            sb.insert(0, '0');
        }
        return sb.toString();
    }

    private String intToBitBinary(int v, int len) {
        StringBuilder sb = new StringBuilder(Integer.toBinaryString(v));
        while (sb.length() % len != 0) {
            sb.insert(0, '0');
        }
        return sb.toString();
    }
}
